package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel, NaiveBayes, NaiveBayesModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo6TrainModel {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("person")
      .config("spark.sql.shuffle.partitions", 1)
      .getOrCreate()

    import spark.implicits._

    //读取处理好的图片数据
    val imageData: DataFrame = spark
      .read
      .format("libsvm")
      .load("data/image_data")

    //将数据拆分成训练集和测试集
    val Array(train: DataFrame, test: DataFrame) = imageData.randomSplit(Array(0.8, 0.2))


    //选择算法
    val naiveBayes = new LogisticRegression()

    //将训练集带入算法训练模型
    val model: LogisticRegressionModel = naiveBayes.fit(train)

    //使用测试集测试模型的准确率
    val testDF: DataFrame = model.transform(test)

    //评估模型的准确率
    val p: Double = testDF.where($"label" === $"prediction").count().toDouble / testDF.count()

    println(s"准确率：$p")

    //将模型保存到hdfs
    model
      .write
      .overwrite()
      .save("data/image_model")

  }

}
